package com.ddone.demo;

import javolution.text.Text;
import org.apache.hadoop.hive.ql.exec.*;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils;

import java.util.ArrayList;
import java.util.List;

/**
 * @author ddone
 * @date 2024/10/15-23:38
 * 将一个点,按照一定的距离辐射出一个一个圆,并输出这些圆的坐标 总共10个经纬度
 * ADD JAR /path/to/your/GenerateCirclePointsUDF.jar;
 * CREATE TEMPORARY FUNCTION generate_circle_points AS 'your.package.GenerateCirclePointsUDF';
 * SELECT expand_point_to_aoi(116.4, 39.9, 1);
 */
@Description(name = "expand_point_to_aoi", value = "_FUNC_(double lon, double lat, double distanceInKm) - returns 10 points around the circle")
public class ExpandPoint2AoiHiveUdf extends GenericUDF {
    private static final double EARTH_RADIUS = 6371.009;
    // 更加精确的地球半径，单位：公里
    private PrimitiveObjectInspector lonOI;
    private PrimitiveObjectInspector latOI;
    private PrimitiveObjectInspector distanceOI;

    // 将角度转换为弧度
    private static double rad(double d) {
        return d * Math.PI / 180.0;
    }

    // 将弧度转换为角度
    private static double deg(double rad) {
        return rad * 180.0 / Math.PI;
    }

    @Override
    public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
        if (arguments.length != 3) {
            throw new UDFArgumentLengthException("The function generate_circle_points requires 3 arguments.");
        }

        // 检查输入参数类型
        if (!(arguments[0] instanceof PrimitiveObjectInspector) ||
                !(arguments[1] instanceof PrimitiveObjectInspector) ||
                !(arguments[2] instanceof PrimitiveObjectInspector)) {
            throw new UDFArgumentTypeException(0, "Only primitive types are accepted as input.");
        }

        lonOI = (PrimitiveObjectInspector) arguments[0];
        latOI = (PrimitiveObjectInspector) arguments[1];
        distanceOI = (PrimitiveObjectInspector) arguments[2];

        // 输出为 Array<String> 类型的对象检查器
        return ObjectInspectorFactory.getStandardListObjectInspector(PrimitiveObjectInspectorFactory.writableStringObjectInspector);
    }

    @Override
    public Object evaluate(DeferredObject[] arguments) throws HiveException {
        double lon = PrimitiveObjectInspectorUtils.getDouble(arguments[0].get(), lonOI);
        double lat = PrimitiveObjectInspectorUtils.getDouble(arguments[1].get(), latOI);
        double distanceInKm = PrimitiveObjectInspectorUtils.getDouble(arguments[2].get(), distanceOI);

        // 生成10个坐标点
        List<Text> result = new ArrayList<>();
        double angularDistance = distanceInKm / EARTH_RADIUS;
        // 角距离（弧度）
        double latRad = rad(lat);
        double lonRad = rad(lon);

        for (int i = 0; i < 10; i++) {
            double angle = 36 * i;
            // 每个点的角度，36度递增
            double angleRad = rad(angle);

            // 计算每个点的纬度
            double latNew = Math.asin(Math.sin(latRad) * Math.cos(angularDistance)
                    + Math.cos(latRad) * Math.sin(angularDistance) * Math.cos(angleRad));

            // 计算每个点的经度
            double lonNew = lonRad + Math.atan2(Math.sin(angleRad) * Math.sin(angularDistance) * Math.cos(latRad),
                    Math.cos(angularDistance) - Math.sin(latRad) * Math.sin(latNew));

            // 将结果转换为字符串，格式为 "lon_lat"
            result.add(new Text(deg(lonNew) + "_" + deg(latNew)));
        }

        return result;
    }

    @Override
    public String getDisplayString(String[] children) {
        return "expand_point_to_aoi(" + children[0] + ", " + children[1] + ", " + children[2] + ")";
    }
}
